#!/usr/bin/env python
# coding=utf-8
# @FileName :main.py
# @Time     :2023/4/10 16:49
# @Author   :ZXC
# @Desc     :数据库连接助手
from urllib import parse

from loguru import logger
from sqlalchemy import text
from sqlalchemy.ext.asyncio import create_async_engine, AsyncSession, async_sessionmaker

from .enums import DatabaseEnum
from .ssh_tunnel import SSHTunnel, SSHTunnelObj
from .schemas import DBObj, DBConn


class DBHelper(object):

    def __init__(self):
        # 数据库连接缓存池
        logger.debug(f'正在初始化数据库连接助手')
        self.connections = dict()

    def get_connection(self, db_obj: DBObj) -> DBConn:
        """
        获取数据库连接
        :param db_obj: 数据库对象模型
        :return: SSH数据库连接对象
        """
        new_db_obj = db_obj.copy()
        # 判断是否连接的是堡垒机自身的数据库
        if new_db_obj.use_ssh and new_db_obj.ssh_host == new_db_obj.host:
            new_db_obj.host = '127.0.0.1'
        # 拼接key
        key = new_db_obj.get_key()
        # 从连接池查找连接对象
        connection = self.connections.get(key, None)
        # 先判断是否已经有connection了，如果有则直接返回。目前还不能判断这个连接是否有效哈
        if connection is not None:
            logger.debug(f'连接池已存在连接对象【{key}】，直接返回{connection}')
            return connection
        # 连接池没有需要的连接对象，重新new一个
        if new_db_obj.use_ssh:
            ssh_tunnel = SSHTunnel(
                SSHTunnelObj(
                    ssh_host=new_db_obj.ssh_host, ssh_port=new_db_obj.ssh_port, ssh_username=new_db_obj.ssh_username,
                    ssh_password=new_db_obj.ssh_password,
                    remote_bind_host=new_db_obj.host, remote_bind_port=new_db_obj.port
                )
            )
            local_port = str(ssh_tunnel.server.local_bind_port)  # 获取转发端口
            # 判断是否启用了ssh
            new_db_obj.host = '127.0.0.1'
            new_db_obj.port = int(local_port)

        # 获取sqlalchemy需要的jdbc url
        jdbc_url = DBHelper.get_jdbc_url(new_db_obj)
        logger.debug(f'获取到数据库连接url【{jdbc_url}】')
        logger.debug(f'准备创建数据库引擎')
        # 创建引擎
        eg = create_async_engine(
            jdbc_url,
            echo=False,  # 如果为True，则将SQLAlchemy的日志输出到标准输出流
            # pool_pre_ping=True,  # 每次执行sql前 悲观检查db是否可用;虽然资源稍微额外的消耗,但是简单可靠
            pool_recycle=1500,  # 用于控制连接池的回收间隔，以秒为单位。应该比mysql中设置的 interactive_timeout 和 wait_timeout 值小 才有效;
            pool_size=5,  # 连接池的大小
            max_overflow=0,  # 允许在连接池被耗尽时，池中可以创建的最大连接数
            pool_timeout=10,  # 获取连接时的超时时间，以秒为单位。应该比 mysql的 connect_timeout值 小或者等于;
        )
        logger.debug(f'准备创建数据库会话')
        ss = async_sessionmaker(bind=eg, class_=AsyncSession)
        # 将数据缓存起来
        data = DBConn(engine=eg, session=ss)
        self.connections[key] = data
        logger.debug(f'添加数据库连接信息到连接池：{data.dict()}')
        return data

    @staticmethod
    def get_jdbc_url(db_obj: DBObj) -> str:
        """
        获取数据库jdbc连接url
        当使用了ssh时，采用local_port端口
        :param db_obj: 数据库对象模型
        :return:
        """
        logger.debug(f'正在获取数据库连接url')
        new_password = parse.quote_plus(db_obj.password)
        if db_obj.sql_type == DatabaseEnum.MYSQL:
            # mysql模式
            return f'mysql+aiomysql://{db_obj.username}:{new_password}@{db_obj.host}:{db_obj.port}/{db_obj.database}'
        if db_obj.sql_type == DatabaseEnum.POSTGRESQL:
            return f'postgresql+aiomysql://{db_obj.username}:{new_password}@{db_obj.host}:{db_obj.port}/{db_obj.database}'
        if db_obj.sql_type == DatabaseEnum.TIDB:
            return f'mysql+aiomysql://{db_obj.username}:{new_password}@{db_obj.host}:{db_obj.port}/{db_obj.database}'
        logger.debug(f'未知的数据库类型：{db_obj.sql_type}')
        raise Exception("未知的数据库类型")

    def remove_connection(self, db_obj: DBObj):
        """
        移除数据库连接
        :param db_obj: 数据库对象模型
        :return:
        """
        # 拼接key
        key = db_obj.get_key()
        logger.debug(f'移除数据库连接:{key}')
        if self.connections.get(key):
            self.connections.pop(key)

    @staticmethod
    async def test_connection(ss: async_sessionmaker):
        """
        测试数据库连接
        :param ss: 数据库会话对象
        :return:
        """
        if ss is None:
            raise Exception("暂不支持的数据库类型")
        async with ss() as session:
            d = await session.execute(text("select 1"))
            logger.info(f'数据库版本为【{d.scalars().fetchall()[0]}】')
            return d
